package main

import (
	"fmt"
	. "go_data_structures_and_algorithms/tool"
)

/*
3. 分隔重排链表
题意：
将一个链表，根据传入的数V,将链表节点值分为小于V,等于V,大于V三部分
再按照小于V,等于V,大于V的顺序，将链表重新组合成新的链表
*/

func main() {
	node1 := &ListNode{
		2,
		nil,
	}
	node2 := &ListNode{
		2,
		nil,
	}
	node3 := &ListNode{
		2,
		nil,
	}
	node4 := &ListNode{
		2,
		nil,
	}
	node5 := &ListNode{
		2,
		nil,
	}
	node6 := &ListNode{
		2,
		nil,
	}

	node1.Next = node2
	node2.Next = node3
	node3.Next = node4
	node4.Next = node5
	node5.Next = node6
	node := listPartition2(node1, 2)
	for node != nil {
		fmt.Println(node.Val)
		node = node.Next
	}

}

func listPartition2(node *ListNode, pivot int) *ListNode {

	if node == nil {
		return nil
	}

	if node.Next == nil {
		return node
	}
	var littleHead *ListNode = nil
	var littleEnd *ListNode = nil
	var equalHead *ListNode = nil
	var equalEnd *ListNode = nil
	var bigHead *ListNode = nil
	var bigEnd *ListNode = nil

	for node != nil {
		if node.Val < pivot {
			if littleEnd == nil && littleHead == nil {
				littleEnd = node
				littleHead = node
				node = node.Next
				littleHead.Next = nil
			} else if littleEnd == littleHead {
				littleEnd = node

				littleHead.Next = littleEnd
				node = node.Next
				littleEnd.Next = nil
			} else {
				littleEnd.Next = node
				littleEnd = littleEnd.Next
				node = node.Next
				littleEnd.Next = nil
			}
		} else if node.Val == pivot {
			if equalHead == nil && equalEnd == nil {
				equalEnd = node
				equalHead = node
				node = node.Next
				equalHead.Next = nil
			} else if equalEnd == equalHead {
				equalEnd = node
				equalHead.Next = equalEnd
				node = node.Next
				equalEnd.Next = nil
			} else {
				equalEnd.Next = node
				equalEnd = equalEnd.Next
				node = node.Next
				equalEnd.Next = nil
			}
		} else if node.Val > pivot {
			if bigHead == nil && bigEnd == nil {
				bigEnd = node
				bigHead = node
				node = node.Next
				bigHead.Next = nil
			} else if bigEnd == bigHead {
				bigEnd = node
				bigHead.Next = bigEnd
				node = node.Next
				bigEnd.Next = nil
			} else {
				bigEnd.Next = node
				bigEnd = bigEnd.Next
				node = node.Next
				bigEnd.Next = nil
			}
		}
	}
	if littleHead != nil {
		littleEnd.Next = equalHead

		if equalHead == nil {

			equalHead = littleEnd
		}
	}

	if equalHead != nil {
		equalEnd.Next = bigHead
		if bigHead == nil {
			bigHead = equalEnd
		}
	}
	if littleHead != nil {
		return littleHead
	} else if equalHead != nil {
		return equalHead
	} else {
		return bigHead
	}
}
